from zipfile import ZipFile
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from IPython.display import Markdown
from typing import List
from collections import Counter

import os
import sys
sys.path.append('../')
from vibrodiagnostics import mafaulda, selection

MAFAULDA_PATH = '../../datasets/MAFAULDA.zip'
FEATURES_PATH =  '../../datasets/features_data/'
MAFAULDA_METADATA = os.path.join(FEATURES_PATH, selection.MAFAULDA_METADATA)

SHAFT_FAULTS = {'normal': 'N', 'imbalance': 'I', 'horizontal-misalignment': 'HM', 'vertical-misalignment': 'VM'}
BEARING_FAULTS = {'overhang-cage_fault': 'O-Cage', 'underhang-cage_fault': 'U-Cage',
                  'underhang-ball_fault': 'U-Ball', 'overhang-ball_fault': 'O-Ball',
                  'underhang-outer_race': 'U-Race', 'overhang-ball_fault': 'O-Race'}
PLACES = [['ax', 'ay', 'az'], ['bx', 'by', 'bz']]


#FAULT_CLASSES = SHAFT_FAULTS  
FAULT_CLASSES = BEARING_FAULTS

NFFT = 2**14
Fs = mafaulda.FS_HZ

files = pd.read_csv(MAFAULDA_METADATA, index_col='filename')
files.reset_index(inplace=True)

Low severity, 2900 RPM

files = files[files['fault'].isin(tuple(FAULT_CLASSES))]
waveform = files[files['rpm'] >= 2900].groupby(by='fault').first()
waveform

	filename	severity	seq	length	duration	rpm	fs
fault
overhang-ball_fault	overhang/ball_fault/0g/49.5616.csv	0g	49.5616.csv	250000	4.99998	2918.323130	50000
overhang-cage_fault	overhang/cage_fault/0g/50.5856.csv	0g	50.5856.csv	250000	4.99998	2965.820725	50000
underhang-ball_fault	underhang/ball_fault/0g/49.5616.csv	0g	49.5616.csv	250000	4.99998	2902.355605	50000
underhang-cage_fault	underhang/cage_fault/0g/49.7664.csv	0g	49.7664.csv	250000	4.99998	2917.865698	50000
underhang-outer_race	underhang/outer_race/0g/49.7664.csv	0g	49.7664.csv	250000	4.99998	2916.459113	50000

def time_waveform(files: List[str], time=None, fs=None):
    for filename in files:
        ts = mafaulda.csv_import(ZipFile(MAFAULDA_PATH), filename)
        display(Markdown(f'**{filename}**'))
        
        for axis in PLACES:
            if time is not None:
                start = int(time * fs)
                duration = fs // 10
                dst = ts[axis].iloc[start:start + duration]
            else:
                dst = ts[axis]

            ax = dst.plot(figsize=(20, 8), grid=True, subplots=True)

            for i, axname in enumerate(axis):
                ax[i].set_xlabel('Time [s]')
                ax[i].set_ylabel(f'Amplitude ({axname}) [m/s^2]')
            plt.show() 


def time_waveform_histogram(files: List[str]):
    for filename in files:
        ts = mafaulda.csv_import(ZipFile(MAFAULDA_PATH), filename)
        display(Markdown(f'**{filename}**'))
        
        axis = PLACES[0] + PLACES[1]
        ts[axis].plot.hist(
            figsize=(15, 4), layout=(2, len(axis) // 2), 
            grid=True, subplots=True, 
            bins=100, xlabel='Amplitude'
        ) #, legend=False)
        plt.tight_layout()
        plt.show()


def plot_class_imbalance(labels):
    counter = Counter(labels)
    for k, v in counter.items():
        per = v / len(labels) * 100
        print(f'Class={k}, n={v} ({per:.3f}%%)')
    
    plt.bar(counter.keys(), counter.values())
    plt.grid(True)
    plt.xticks(rotation=45)
    plt.show()


def spectogram(x):
    fig, ax = plt.subplots(figsize=(15, 4))
    cmap = plt.get_cmap('inferno')
    pxx, freqs, t, im = plt.specgram(
        x, NFFT=NFFT, Fs=Fs,
        detrend='mean',
        mode='magnitude', scale='dB',
        cmap=cmap, vmin=-60
    )
    fig.colorbar(im, aspect=20, pad=0.04)
    ax.set_xlabel('Time [s]')
    ax.set_ylabel('Frequency [Hz]')
    mafaulda.resolution_calc(Fs, NFFT)
    return freqs, pxx


def window_idx(t):
    return (Fs * t) // NFFT + 1


def spectrum_slice(freqs, pxx, t):
    fig, ax = plt.subplots(2, 1, figsize=(20, 8))
    n = window_idx(t)

    dB = 20 * np.log10(pxx.T[n] / 0.000001)
    ax[0].plot(freqs, dB)      # 1 dB = 1 um/s^2
    ax[0].grid(True)
    ax[0].set_xlabel('Frequency [Hz]')
    ax[0].set_ylabel('Amplitude [dB]')
    
    ax[1].plot(freqs, pxx.T[n])
    ax[1].grid(True)
    ax[1].set_xlabel('Frequency [Hz]')
    ax[1].set_ylabel('Amplitude [m/s^2]')
    return n


def get_spectrograms(files: List[str], axis: str) -> list:
    spectrograms = []

    for filename in files:
        ts = mafaulda.csv_import(ZipFile(MAFAULDA_PATH), filename)
        display(Markdown(f'**{filename}** ({axis.upper()})'))
        freqs, pxx = spectogram(ts[axis])
        spectrograms.append((filename, freqs, pxx))
        plt.show()

    return spectrograms  


def show_spectrogram_detail(spectrograms: list, axis: str, t: float):
    for name, freqs, Pxx in spectrograms:
        display(Markdown(f'**{name}** ({axis.upper()}) axis @ {t}s)'))
        i_window = spectrum_slice(freqs, Pxx, t)
        plt.show() 

plot_class_imbalance(files['fault'])

Class=overhang-ball_fault, n=137 (15.515%%)
Class=overhang-cage_fault, n=188 (21.291%%)
Class=underhang-ball_fault, n=186 (21.065%%)
Class=underhang-cage_fault, n=188 (21.291%%)
Class=underhang-outer_race, n=184 (20.838%%)

time_waveform(waveform['filename'])

overhang/ball_fault/0g/49.5616.csv

overhang/cage_fault/0g/50.5856.csv

underhang/ball_fault/0g/49.5616.csv

underhang/cage_fault/0g/49.7664.csv

underhang/outer_race/0g/49.7664.csv

time_waveform(waveform['filename'], time=2, fs=mafaulda.FS_HZ)

overhang/ball_fault/0g/49.5616.csv

overhang/cage_fault/0g/50.5856.csv

underhang/ball_fault/0g/49.5616.csv

underhang/cage_fault/0g/49.7664.csv

underhang/outer_race/0g/49.7664.csv

time_waveform_histogram(waveform['filename'])

overhang/ball_fault/0g/49.5616.csv

overhang/cage_fault/0g/50.5856.csv

underhang/ball_fault/0g/49.5616.csv

underhang/cage_fault/0g/49.7664.csv

underhang/outer_race/0g/49.7664.csv

A measurement point

ax_spectra = get_spectrograms(waveform['filename'], 'ax')

overhang/ball_fault/0g/49.5616.csv (AX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

overhang/cage_fault/0g/50.5856.csv (AX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/ball_fault/0g/49.5616.csv (AX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/cage_fault/0g/49.7664.csv (AX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/outer_race/0g/49.7664.csv (AX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

show_spectrogram_detail(ax_spectra, 'ax', t=2)

overhang/ball_fault/0g/49.5616.csv (AX) axis @ 2s)

overhang/cage_fault/0g/50.5856.csv (AX) axis @ 2s)

underhang/ball_fault/0g/49.5616.csv (AX) axis @ 2s)

underhang/cage_fault/0g/49.7664.csv (AX) axis @ 2s)

underhang/outer_race/0g/49.7664.csv (AX) axis @ 2s)

ay_spectra = get_spectrograms(waveform['filename'], 'ay')

overhang/ball_fault/0g/49.5616.csv (AY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

overhang/cage_fault/0g/50.5856.csv (AY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/ball_fault/0g/49.5616.csv (AY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/cage_fault/0g/49.7664.csv (AY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/outer_race/0g/49.7664.csv (AY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

show_spectrogram_detail(ay_spectra, 'ay', t=2)

overhang/ball_fault/0g/49.5616.csv (AY) axis @ 2s)

overhang/cage_fault/0g/50.5856.csv (AY) axis @ 2s)

underhang/ball_fault/0g/49.5616.csv (AY) axis @ 2s)

underhang/cage_fault/0g/49.7664.csv (AY) axis @ 2s)

underhang/outer_race/0g/49.7664.csv (AY) axis @ 2s)

az_spectra = get_spectrograms(waveform['filename'], 'az')

overhang/ball_fault/0g/49.5616.csv (AZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

overhang/cage_fault/0g/50.5856.csv (AZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/ball_fault/0g/49.5616.csv (AZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/cage_fault/0g/49.7664.csv (AZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/outer_race/0g/49.7664.csv (AZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

show_spectrogram_detail(az_spectra, 'az', t=2)

overhang/ball_fault/0g/49.5616.csv (AZ) axis @ 2s)

overhang/cage_fault/0g/50.5856.csv (AZ) axis @ 2s)

underhang/ball_fault/0g/49.5616.csv (AZ) axis @ 2s)

underhang/cage_fault/0g/49.7664.csv (AZ) axis @ 2s)

underhang/outer_race/0g/49.7664.csv (AZ) axis @ 2s)

B measurement point

bx_spectra = get_spectrograms(waveform['filename'], 'bx')

overhang/ball_fault/0g/49.5616.csv (BX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

overhang/cage_fault/0g/50.5856.csv (BX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/ball_fault/0g/49.5616.csv (BX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/cage_fault/0g/49.7664.csv (BX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/outer_race/0g/49.7664.csv (BX)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

show_spectrogram_detail(bx_spectra, 'bx', t=2)

overhang/ball_fault/0g/49.5616.csv (BX) axis @ 2s)

overhang/cage_fault/0g/50.5856.csv (BX) axis @ 2s)

underhang/ball_fault/0g/49.5616.csv (BX) axis @ 2s)

underhang/cage_fault/0g/49.7664.csv (BX) axis @ 2s)

underhang/outer_race/0g/49.7664.csv (BX) axis @ 2s)

by_spectra = get_spectrograms(waveform['filename'], 'by')

overhang/ball_fault/0g/49.5616.csv (BY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

overhang/cage_fault/0g/50.5856.csv (BY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/ball_fault/0g/49.5616.csv (BY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/cage_fault/0g/49.7664.csv (BY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/outer_race/0g/49.7664.csv (BY)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

show_spectrogram_detail(by_spectra, 'by', t=2)

overhang/ball_fault/0g/49.5616.csv (BY) axis @ 2s)

overhang/cage_fault/0g/50.5856.csv (BY) axis @ 2s)

underhang/ball_fault/0g/49.5616.csv (BY) axis @ 2s)

underhang/cage_fault/0g/49.7664.csv (BY) axis @ 2s)

underhang/outer_race/0g/49.7664.csv (BY) axis @ 2s)

bz_spectra = get_spectrograms(waveform['filename'], 'bz')

overhang/ball_fault/0g/49.5616.csv (BZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

overhang/cage_fault/0g/50.5856.csv (BZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/ball_fault/0g/49.5616.csv (BZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/cage_fault/0g/49.7664.csv (BZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

underhang/outer_race/0g/49.7664.csv (BZ)

Window size: 16384
Heinsenberg box
	Time step: 327.68 ms
	Frequency step: 3.0517578125 Hz

show_spectrogram_detail(bz_spectra, 'bz', t=2)

overhang/ball_fault/0g/49.5616.csv (BZ) axis @ 2s)

overhang/cage_fault/0g/50.5856.csv (BZ) axis @ 2s)

underhang/ball_fault/0g/49.5616.csv (BZ) axis @ 2s)

underhang/cage_fault/0g/49.7664.csv (BZ) axis @ 2s)

underhang/outer_race/0g/49.7664.csv (BZ) axis @ 2s)

Features

# box plots